from selenium import webdriver
import re
from urllib.parse import unquote


# 自定义异常，用于在找不到源直播推送流时引发
class LiveStreamNotFoundException(Exception):
    pass


class LiveSource(object):
    headers = {
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)"
        " Chrome/98.0.4758.82 Safari/537.36"
    }

    # 初始化web 浏览器
    def init_webdrive(self):
        chrome_option = webdriver.ChromeOptions()
        chrome_option.add_experimental_option("excludeSwitches", ["enable-automation"])
        chrome_option.add_argument("--no-sandbox")
        chrome_option.add_argument("--disable-dev-shm-usage")
        chrome_option.add_argument("--disable-blink-features=AutomationControlled")
        chrome_option.add_argument("start-maximized")
        chrome_option.add_argument(f'user-agent={self.headers["user-agent"]}')
        chrome_option.add_argument("headless")  # 静默模式
        # 扫码登录
        # chrome_option.add_argument(r'--user-data-dir=ChromeUserData')
        # time.sleep(25)
        # s = Service(r"./chromedriver.exe")
        driver = webdriver.Chrome(options=chrome_option)
        driver.execute_cdp_cmd(
            "Page.addScriptToEvaluateOnNewDocument",
            {
                "source": """
                                     Object.defineProperty(navigator, 'webdriver', {
                                       get: () => undefined
                                     })
                                   """
            },
        )
        return driver

    def async_url(self, live_url):
        try:
            driver = self.init_webdrive()
            driver.get(live_url)
            match_text = driver.page_source.encode().decode("unicode_escape")
            live_stream = re.findall(r"\"hls_pull_url\":\"(.*?)\"", unquote(match_text))
            driver.close()

            if live_stream:
                # 如果找到了直播流地址，返回第一个元素
                return live_stream[0]
            else:
                # 如果未找到直播流地址，引发自定义异常
                raise LiveStreamNotFoundException("找不到源直播推送流")
        except Exception as e:
            # 捕获所有异常并打印错误信息
            print(f"发生异常: {str(e)}")
            return None
